#include <iostream>
#include <cstdio>
#include <cstdlib>

#include <chTimer.hpp>

#define GPUS 1

__global__ void saxpy_parallel(int n, float alpha, float *x, float *y)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;

	if (i<n)
	{
		y[i] = alpha*x[i] + y[i];
	}
}

extern "C" void do_saxpy_on_cpu(float *x, float *y, float alpha, int size)
{
	for (int i=0; i<size; i++)
	{
		y[i] = alpha*x[i] + y[i];
	}
}

extern "C" void do_saxpy_on_gpu(float *global_x, float *global_y, float alpha, int numThreads, int numBlocks, size_t N, double *results)
{
	ChTimer h2d_timer, d2h_timer, kernel_timer;

	#if 0
    // Number of CUDA devices
    int devCount;
    cudaGetDeviceCount(&devCount);
    printf("There are %d CUDA devices.\n", devCount);
	#endif

#if GPUS == 2

	/* device memory */
	float *d_x[2];
	float *d_y[2];
	cudaSetDevice(0);
	cudaMalloc((void**)&d_x[0], (size_t)((size/2) * sizeof(float)));
	cudaMalloc((void**)&d_y[0], (size_t)((size/2) * sizeof(float)));
	cudaSetDevice(1);
	cudaMalloc((void**)&d_x[1], (size_t)((size/2) * sizeof(float)));
	cudaMalloc((void**)&d_y[1], (size_t)((size/2) * sizeof(float)));

	/* copy data to device */
	h2d_timer.start();
	cudaSetDevice(0);
	cudaMemcpy(&d_x[0], &global_x[0], (size/2) * sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(&d_y[0], &global_y[0], (size/2) * sizeof(float), cudaMemcpyHostToDevice);
	cudaSetDevice(1);
	cudaMemcpy(&d_x[1], &global_x[size/2], (size/2) * sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(&d_y[1], &global_y[size/2], (size/2) * sizeof(float), cudaMemcpyHostToDevice);
	h2d_timer.stop();

	/* start kernel */
	kernel_timer.start();
	cudaSetDevice(0);
	saxpy_parallel<<<numBlocks, numThreads>>>(size/2, alpha, &d_x[0][0], &d_y[0][0]);
	cudaSetDevice(1);
	saxpy_parallel<<<numBlocks, numThreads>>>(size/2, alpha, &d_x[1][0], &d_y[1][0]);
	cudaThreadSynchronize();
	kernel_timer.stop();

	/* read data from device */
	d2h_timer.start();
	cudaSetDevice(0);
	cudaMemcpy(&global_y[0],      &d_y[0], (size/2) * sizeof(float), cudaMemcpyDeviceToHost);
	cudaSetDevice(1);
	cudaMemcpy(&global_y[size/2], &d_y[1], (size/2) * sizeof(float), cudaMemcpyDeviceToHost);
	d2h_timer.stop();

	/* clean up*/
	cudaFree(&d_x[0][0]);
	cudaFree(&d_y[0][0]);
	cudaFree(&d_x[1][0]);
	cudaFree(&d_y[1][0]);

#elif GPUS == 1

	/* device memory */
	float *d_x;
	float *d_y;
	cudaMalloc((void**)&d_x, N*sizeof(float));
	cudaMalloc((void**)&d_y, N*sizeof(float));

	/* copy data to device */
	h2d_timer.start();
	cudaMemcpy(d_x, global_x, N*sizeof(float), cudaMemcpyHostToDevice);
	cudaMemcpy(d_y, global_y, N*sizeof(float), cudaMemcpyHostToDevice);
	h2d_timer.stop();

	/* start kernel */
	kernel_timer.start();
	saxpy_parallel<<<numBlocks, numThreads>>>(N, alpha, d_x, d_y);
	cudaThreadSynchronize();
	kernel_timer.stop();

	/* read data from device */
	d2h_timer.start();
	cudaMemcpy(global_y, d_y, N*sizeof(float), cudaMemcpyDeviceToHost);
	d2h_timer.stop();

	/* clean up*/
	cudaFree(d_x);
	cudaFree(d_y);

#else
	#error invalid gpu number
#endif
	
	double bytes_h2d = (double)(2*N*sizeof(float));
	double bytes_d2h = (double)(N*sizeof(float));

	results[0] = h2d_timer.getTime();
	results[1] = h2d_timer.getBandwidth(bytes_h2d);
	results[2] = d2h_timer.getTime();
	results[3] = d2h_timer.getBandwidth(bytes_d2h);
	results[4] = kernel_timer.getTime();
}

